/*
 * Copyright(C) 2022. countabc Technologies Co.,Ltd. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "MediaDeal_V2.h"

#include "MxBase/Log/Log.h"
#include <algorithm>
#include <map>
#include "MxBase/Maths/FastMath.h"
#include <condition_variable>
#include <signal.h> // signal functions
#include "cropResizePaste.hpp"
#define DEBUG
using namespace MxBase;
using namespace std;
namespace
{
    int countsave = 0;
    const uint32_t TIME_POP = 1000;
    const uint32_t SRC_WIDTH = 1280;
    const uint32_t SRC_HEIGHT = 720;
    const uint32_t MAX_WIDTH = 3840;
    const uint32_t MAX_HEIGHT = 2160;
    const uint32_t RESIZE_WIDTH = 1088;
    const uint32_t RESIZE_HEIGHT = 608;
    const uint32_t MAX_FRAME_COUNT = 5000;
    const uint32_t TARGET_FRAME_COUNT = 900;
    const uint32_t TARGET_FRAME = TARGET_FRAME_COUNT - 5;
    const uint32_t MS_TIMEOUT = 2000;
    const uint32_t DEVICE_ID = 0;
    const uint32_t CHANNEL_ID = 1;
    AVFormatContext *pFormatCtx = nullptr;
    const uint32_t SECOND = 40;
    const uint32_t FIFTY = 50;
    const uint32_t RATE = 25;
    const std::string modelPath = "../models/mot_v2.om";
    const std::string savePath = "../saveimage/";
    V2Param v2Param(DEVICE_ID, modelPath);
    std::shared_ptr<ascendVehicleTracking::MOTConnection> tracker = std::make_shared<ascendVehicleTracking::MOTConnection>();
    std::shared_ptr<FairMOT> fairModel = std::make_shared<FairMOT>(v2Param);
}

APP_ERROR MediaDeal_V2::Init(std::string filePath, std::string savePath)
{
    openFilePath = filePath;
    saveFilePath = savePath;

    PullStream(filePath);

    return APP_ERR_OK;
}

// pull stream
void MediaDeal_V2::PullStream(std::string filePath)
{
    LogInfo << "start to PullStream";
    av_register_all();
    avformat_network_init();
    pFormatCtx = avformat_alloc_context();
    pFormatCtx = CreateFormatContext(filePath);
    av_dump_format(pFormatCtx, 0, filePath.c_str(), 0);
}

// ffmpeg 拉流
AVFormatContext *MediaDeal_V2::CreateFormatContext(std::string filePath)
{
    LogInfo << "start to CreatFormatContext!";
    // creat message for stream pull
    AVFormatContext *formatContext = nullptr;
    AVDictionary *options = nullptr;

    LogInfo << "start to avformat_open_input!";
    int ret = avformat_open_input(&formatContext, filePath.c_str(), nullptr, &options);
    if (options != nullptr)
    {
        av_dict_free(&options);
    }
    if (ret != 0)
    {
        LogError << "Couldn`t open input stream " << filePath.c_str() << " ret=" << ret;
        return nullptr;
    }
    ret = avformat_find_stream_info(formatContext, nullptr);
    if (ret != 0)
    {
        LogError << "Couldn`t open input stream information";
        return nullptr;
    }

    return formatContext;
}

// 获取H264中的帧
void MediaDeal_V2::GetFrame(AVPacket &pkt, FrameImage &frameImage, AVFormatContext *pFormatCtx)
{
    av_init_packet(&pkt);
    int ret = av_read_frame(pFormatCtx, &pkt);
    if (ret != 0)
    {
        LogInfo << "[StreamPuller] channel Read frame failed, continue!";
        if (ret == AVERROR_EOF)
        {
            LogInfo << "[StreamPuller] channel StreamPuller is EOF, over!";
            sleep(1);
            return;
        }
    }
    else
    {
        if (pkt.size <= 0)
        {
            LogInfo << "Invalid pkt.size: " << pkt.size;
            return;
        }
        // send to the device
        auto hostDeleter = [](void *dataPtr) -> void {};
        MemoryData data(pkt.size, MemoryData::MEMORY_HOST);
        MemoryData src((void *)(pkt.data), pkt.size, MemoryData::MEMORY_HOST);
        APP_ERROR ret = MemoryHelper::MxbsMallocAndCopy(data, src);
        if (ret != APP_ERR_OK)
        {
            LogError << "MxbsMallocAndCopy failed!";
        }
        std::shared_ptr<uint8_t> imageData((uint8_t *)data.ptrData, hostDeleter);
        Image subImage(imageData, pkt.size);
        frameImage.image = subImage;
#ifdef USE_200DK
        subImage.ToDevice(0);
        frameImage.image = subImage;
        frameImage.deviceId = 0;
#endif
        LogInfo << "'channelId = " << frameImage.channelId << ", frameId = " << frameImage.frameId << " , dataSize = " << frameImage.image.GetDataSize();

        av_packet_unref(&pkt);
    }
    return;
}

// resize frame image
APP_ERROR MxBase::Resize(const MxBase::Image &decodedImage, MxBase::Image &resizeImage)
{
    std::shared_ptr<MxBase::ImageProcessor> imageProcessorDptr;
    // imageProcess init
    imageProcessorDptr = std::make_shared<MxBase::ImageProcessor>(DEVICE_ID);
    if (imageProcessorDptr == nullptr)
    {
        LogError << "imageProcessorDptr nullptr";
    }

    APP_ERROR ret;

    // set size param
    Size resizeConfig(RESIZE_WIDTH, RESIZE_HEIGHT);
    resizeImage = resizeKeepAspectRatioFit(decodedImage.GetSize().width, decodedImage.GetSize().height,
                                           RESIZE_WIDTH, RESIZE_HEIGHT, decodedImage, *imageProcessorDptr);

    return APP_ERR_OK;
}

// pull stream thread
void MediaDeal_V2::PullStreamThread()
{
    LogInfo << "start to pull stream thread.";
    AVPacket pkt;
    uint32_t frameId = 0;
    MxBase::DeviceContext device;
    device.devId = DEVICE_ID;
    APP_ERROR ret = MxBase::DeviceManager::GetInstance()->SetDevice(device);
    if (ret != APP_ERR_OK)
    {
        LogError << "SetDevice failed";
        return;
    }

    while (!stopFlag)
    {
        std::this_thread::sleep_for(std::chrono::milliseconds(SECOND));
        Image subImage;
        FrameImage frame;
        frame.channelId = 0;
        frame.frameId = frameId;
        frame.image = subImage;
        MediaDeal_V2::GetFrame(pkt, frame, pFormatCtx);

        pullStreamQueue.Push(frame);
    }
    return;
}

void MediaDeal_V2::DecodeThread()
{
    LogInfo << "start to video decode thread.";
    int32_t channelId = CHANNEL_ID;
    // 解码器参数
    VideoDecodeConfig vDconfig;
    VideoDecodeCallBack cPtr = MxBase::CallBackVdecode;
    vDconfig.width = SRC_WIDTH;
    vDconfig.height = SRC_HEIGHT;
    vDconfig.callbackFunc = cPtr;
    vDconfig.skipInterval = 0; // 跳帧控制

    std::shared_ptr<VideoDecoder> videoDecoder = std::make_shared<VideoDecoder>(vDconfig, DEVICE_ID, channelId);
    while (!stopFlag)
    {
        FrameImage temp;
        APP_ERROR ret = pullStreamQueue.Pop(temp, TIME_POP);
        if (ret != APP_ERR_OK)
        {
            LogError << "pullStreamQueue pop failed" << ret;
            continue;
        }

        ret = videoDecoder->Decode(temp.image.GetData(), temp.image.GetDataSize(), frameId, &decodeQueue);
        if (ret != APP_ERR_OK)
        {
            LogError << "videoDecoder Decode failed. ret is: " << ret;
        }
        frameId += 1;
    }
    return;
}

// 视频解码回调
APP_ERROR MxBase::CallBackVdecode(Image &decodedImage, uint32_t channelId, uint32_t frameId, void *userData)
{
    FrameImage frameImage;
    frameImage.image = decodedImage;
    frameImage.channelId = channelId;

    BlockingQueue<FrameImage> *p_decodeQueue = (BlockingQueue<FrameImage> *)userData;
    p_decodeQueue->Push(frameImage);

    return APP_ERR_OK;
};

// resize video frame
void MediaDeal_V2::ResizeThread()
{
    LogInfo << "start to resize resize thread.";
    while (!stopFlag)
    {
        FrameImage decodeTemp;
        APP_ERROR ret = decodeQueue.Pop(decodeTemp, TIME_POP);
        if (ret != APP_ERR_OK)
        {
            LogError << "decodeQueue pop failed" << ret;
            continue;
        }

        MxBase::Image resizeImage;
        Resize(decodeTemp.image, resizeImage);
// the code to save test code
#ifdef DEBUG
        countsave++;
        string filename = savePath + to_string(countsave) + "temp.jpg";
        std::shared_ptr<MxBase::ImageProcessor> imageProcessorDptrre;
        // imageProcess init
        imageProcessorDptrre = std::make_shared<MxBase::ImageProcessor>(DEVICE_ID);
        if (imageProcessorDptrre == nullptr)
        {
            LogError << "imageProcessorDptrre nullptr";
        }
        imageProcessorDptrre->Encode(decodeTemp.image, filename);
#endif
#ifndef USE_200DK
        decodeTemp.image.ToHost();
#endif
        // backup image
        backupQueue.Push(decodeTemp);
        FrameImage frameImage;
        frameImage.image = resizeImage;
        frameImage.channelId = channelId;

        resizeQueue2infer.Push(frameImage);
    }
    return;
}

// infer frameimage
void MediaDeal_V2::inferThread()
{
    LogInfo << "start to infer thread.";
    APP_ERROR ret;
    while (!stopFlag)
    {
        FrameImage resizeTemp;
        ret = resizeQueue2infer.Pop(resizeTemp, TIME_POP);
        if (ret != APP_ERR_OK)
        {
            LogError << "resizeQueue2infer pop failed" << ret;
            continue;
        }
        std::vector<MxBase::Tensor> outputs;
        ret = fairModel->Infer(resizeTemp.image, outputs);
        if (ret != APP_ERR_OK)
        {
            LogError << "infer failed. ret is: " << ret;
        }

        inferQueue2pocessone.Push(outputs);
    }
    return;
}

// postone process
void MediaDeal_V2::PostoneThread()
{
    LogInfo << "start to Postone thread.";
    APP_ERROR ret;
    while (!stopFlag)
    {
        std::vector<std::vector<MxBase::ObjectInfo>> objectInfos;

        std::vector<MxBase::Tensor> outputs;
        ret = inferQueue2pocessone.Pop(outputs, TIME_POP);
        if (ret != APP_ERR_OK)
        {
            LogError << "inferQueue2pocessone pop failed" << ret;
            continue;
        }

        ret = fairModel->PostProcessone(outputs, objectInfos);
        if (ret != APP_ERR_OK)
        {
            LogError << "PostProcessone failed. ret is: " << ret;
        }

        postpocessoneQueue2track.Push(objectInfos);
    }
    return;
}

// track process
void MediaDeal_V2::TrackThread()
{
    LogInfo << "start to track thread.";
    APP_ERROR ret;
    while (!stopFlag)
    {
        std::vector<MxBase::ObjectInfo> objectInfos;

        std::vector<std::vector<MxBase::ObjectInfo>> postoneoutputs;
        ret = postpocessoneQueue2track.Pop(postoneoutputs, TIME_POP);
        if (ret != APP_ERR_OK)
        {
            LogError << "postpocessoneQueue2track pop failed" << ret;
            continue;
        }

        ret = tracker->ProcessSort(postoneoutputs);
        if (ret != APP_ERR_OK)
        {
            LogError << "ProcessSort failed. ret is: " << ret;
        }
        ret = tracker->GettrackResult(objectInfos);
        if (ret != APP_ERR_OK)
        {
            LogError << "No tracker";
        }
        trackQueue2pocesstwo.Push(objectInfos);
    }
    return;
}

// posttwo process
void MediaDeal_V2::PosttwoThread()
{
    LogInfo << "start to Posttwo thread.";
    APP_ERROR ret;
    while (!stopFlag)
    {
        FrameImage Posttwoimage;
        ret = backupQueue.Pop(Posttwoimage, TIME_POP);
        if (ret != APP_ERR_OK)
        {
            LogError << "backupQueue pop failed" << ret;
            continue;
        }

        std::vector<MxBase::ObjectInfo> outputs;
        ret = trackQueue2pocesstwo.Pop(outputs, TIME_POP);
        if (ret != APP_ERR_OK)
        {
            LogError << "trackQueue2pocesstwo pop failed" << ret;
            continue;
        }

        ret = fairModel->PostProcesstwo(outputs, Posttwoimage);
        if (ret != APP_ERR_OK)
        {
            LogError << "PostProcesstwo failed. ret is: " << ret;
        }
        postpocesstwo2encodeQueue.Push(Posttwoimage);
    }
    return;
}

void MediaDeal_V2::EncodeThread()
{
    LogInfo << "start to video encode thread.";
    VideoEncodeConfig vEConfig;
    VideoEncodeCallBack cEPtr = MxBase::CallBackVencode;
    vEConfig.callbackFunc = cEPtr;
    vEConfig.width = SRC_WIDTH;
    vEConfig.height = SRC_HEIGHT;
    // 用户可自定义编码参数
    vEConfig.keyFrameInterval = FIFTY;
    vEConfig.srcRate = RATE;
    vEConfig.rcMode = 0;
    vEConfig.maxBitRate = MS_TIMEOUT;
    vEConfig.ipProp = FIFTY;
    std::shared_ptr<VideoEncoder> videoEncoder = std::make_shared<VideoEncoder>(vEConfig, DEVICE_ID);
    int frameCount = 0;
    while (!stopFlag)
    {
        FrameImage resizeTemp;
        APP_ERROR ret = postpocesstwo2encodeQueue.Pop(resizeTemp, TIME_POP);
        if (ret != APP_ERR_OK)
        {
            LogError << "postpocesstwo2encodeQueue pop failed" << ret;
            continue;
        }
        videoEncoder->Encode(resizeTemp.image, resizeTemp.frameId, &encodeQueue);
        frameCount += 1;
    }
    return;
}

// 视频编码回调
APP_ERROR MxBase::CallBackVencode(std::shared_ptr<uint8_t> &outDataPtr, uint32_t &outDataSize, uint32_t &channelId,
                                  uint32_t &frameId, void *userData)
{
    Image image(outDataPtr, outDataSize, -1, Size(MAX_WIDTH, MAX_HEIGHT));
    FrameImage frameImage;
    frameImage.image = image;
    frameImage.channelId = channelId;
    frameImage.frameId = frameId;

    LogInfo << "frameId(" << frameImage.frameId << ") encoded successfully.";

    bool bIsIDR = (outDataSize > 1);
    if (frameImage.frameId)
    {
        if (!bIsIDR)
        {
            LogError << "Not bIsIDR!!!";
            return APP_ERR_OK;
        }
    }

    BlockingQueue<FrameImage> *p_encodeQueue = (BlockingQueue<FrameImage> *)userData;
    p_encodeQueue->Push(frameImage);
    return APP_ERR_OK;
};

// records frame number per second
void MediaDeal_V2::CalFps()
{
    while (!stopFlag)
    {
        sleep(1);
        LogInfo << "frame number: " << finishCount - lastCount;
        LogInfo << "video encode frame rate for per second: " << (finishCount - lastCount) / 1 << " fps.";
        lastCount = finishCount;
    }
}

// write result
void MediaDeal_V2::WriteThread()
{
    FILE *fp = fopen(saveFilePath.c_str(), "wb");
    if (fp == nullptr)
    {
        LogError << "Failed to open file.";
        return;
    }

    while (!stopFlag)
    {
        FrameImage encodeTemp;
        APP_ERROR ret = encodeQueue.Pop(encodeTemp, TIME_POP);
        if (ret != APP_ERR_OK)
        {
            LogError << "encodeQueue pop failed" << ret;
            continue;
        }

        std::shared_ptr<uint8_t> data_sp = encodeTemp.image.GetData();
        void *data_p = data_sp.get();
        if (fwrite(data_p, encodeTemp.image.GetDataSize(), 1, fp) != 1)
        {
            LogInfo << "write frame to file fail";
        }

        finishCount++;
        if (finishCount == TARGET_FRAME)
        {
            stopFlag = true;
        }
    }
    return;
}

// process thread
APP_ERROR MediaDeal_V2::Process(std::string filePath, std::string outPath)
{
    auto startTime = std::chrono::high_resolution_clock::now();

    std::thread threadPullStream(&MediaDeal_V2::PullStreamThread, this);

    std::thread threadDecode(&MediaDeal_V2::DecodeThread, this);

    std::thread threadResize(&MediaDeal_V2::ResizeThread, this);

    std::thread threadInfer(&MediaDeal_V2::inferThread, this);

    std::thread threadPostone(&MediaDeal_V2::PostoneThread, this);

    std::thread threadTrack(&MediaDeal_V2::TrackThread, this);

    std::thread threadPosttwo(&MediaDeal_V2::PosttwoThread, this);

    std::thread threadEncode(&MediaDeal_V2::EncodeThread, this);

    std::thread threadWrite(&MediaDeal_V2::WriteThread, this);

    std::thread threadCalFps(&MediaDeal_V2::CalFps, this);

    threadPullStream.join();
    threadDecode.join();
    threadResize.join();
    threadInfer.join();
    threadPostone.join();
    threadTrack.join();
    threadPosttwo.join();
    threadEncode.join();
    threadWrite.join();
    threadCalFps.join();

    auto endTime = std::chrono::high_resolution_clock::now();
    double costS = std::chrono::duration<double>(endTime - startTime).count();
    LogInfo << "total process time: " << costS << "s.";
    double fps = 0;
    if (costS == 0)
    {
        fps = 1;
    }
    else
    {
        fps = finishCount / costS;
    }
    LogInfo << "Total decode frame rate: " << fps << " fps.";
    return APP_ERR_OK;
}

// stop process
void MediaDeal_V2::stopProcess()
{
    stopFlag = true;
}
